#include <string>
#include <iostream>
#include <unordered_map>
#include <queue>
#include <pthread.h>
#include <mpi.h> // 添加 MPI 头文件
using namespace std;

void send_string(const std::string& s, int dest, int tag, MPI_Comm comm);
std::string recv_string(int source, int tag, MPI_Comm comm, MPI_Status* status);

class segment {
public:
    int type; // 0: 未设置, 1: 字母, 2: 数字, 3: 特殊字符
    int length; // 长度，例如S6的长度就是6
    segment(int type, int length) : type(type), length(length) {}
    segment() : type(0), length(0) {}
    // 打印相关信息（内联实现示例）
    void PrintSeg() {
        cout << "Segment: Type " << type << ", Length " << length << endl;
    }

    vector<string> ordered_values;
    vector<int> ordered_freqs;
    int total_freq = 0;
    unordered_map<string, int> values;
    unordered_map<int, int> freqs;

    void insert(string value); // 声明保持不变
    void order(); // 声明保持不变
    void PrintValues(); // 声明保持不变
};

class PT {
public:
    vector<segment> content;
    int pivot = 0;
    void insert(segment seg); // 声明保持不变
    void PrintPT(); // 声明保持不变
    vector<PT> NewPTs(); // 声明保持不变
    vector<int> curr_indices;
    vector<int> max_indices;
    float preterm_prob;
    float prob;
    PT() = default;
    // 添加拷贝构造函数（可选，确保容器安全）
    PT(const PT& other) = default;
    PT& operator=(const PT& other) = default;
};

class model {
public:
    // 原有成员保持不变
    int preterm_id = -1;
    int letters_id = -1;
    int digits_id = -1;
    int symbols_id = -1;
    int GetNextPretermID() { return ++preterm_id; }
    int GetNextLettersID() { return ++letters_id; }
    int GetNextDigitsID() { return ++digits_id; }
    int GetNextSymbolsID() { return ++symbols_id; }

    int total_preterm = 0;
    vector<PT> preterminals;
    int FindPT(PT pt); // 声明保持不变
    vector<segment> letters;
    vector<segment> digits;
    vector<segment> symbols;
    int FindLetter(segment seg); // 声明保持不变
    int FindDigit(segment seg); // 声明保持不变
    int FindSymbol(segment seg); // 声明保持不变
    unordered_map<int, int> preterm_freq;
    unordered_map<int, int> letters_freq;
    unordered_map<int, int> digits_freq;
    unordered_map<int, int> symbols_freq;
    vector<PT> ordered_pts;

    void train(string train_path); // 声明保持不变
    void store(string store_path); // 声明保持不变
    void load(string load_path); // 声明保持不变
    void parse(string pw); // 声明保持不变
    void order(); // 声明保持不变
    void print(); // 声明保持不变

    // 添加 MPI 广播接口
    void broadcast() {
        // 广播基础数据
        MPI_Bcast(&preterm_id, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&letters_id, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&digits_id, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&symbols_id, 1, MPI_INT, 0, MPI_COMM_WORLD);

        // 广播容器大小（需先发送大小再发送数据）
        int letters_size = letters.size();
        MPI_Bcast(&letters_size, 1, MPI_INT, 0, MPI_COMM_WORLD);
        letters.resize(letters_size);
        if (!letters.empty()) {
            MPI_Bcast(letters.data(), letters_size * sizeof(segment), MPI_BYTE, 0, MPI_COMM_WORLD);
        }

        // 类似处理 digits 和 symbols
        int digits_size = digits.size();
        MPI_Bcast(&digits_size, 1, MPI_INT, 0, MPI_COMM_WORLD);
        digits.resize(digits_size);
        if (!digits.empty()) {
            MPI_Bcast(digits.data(), digits_size * sizeof(segment), MPI_BYTE, 0, MPI_COMM_WORLD);
        }

        int symbols_size = symbols.size();
        MPI_Bcast(&symbols_size, 1, MPI_INT, 0, MPI_COMM_WORLD);
        symbols.resize(symbols_size);
        if (!symbols.empty()) {
            MPI_Bcast(symbols.data(), symbols_size * sizeof(segment), MPI_BYTE, 0, MPI_COMM_WORLD);
        }
    }
};

class PriorityQueue {
public:
    vector<PT> priority;
    vector<string> guesses;
    model m;
    int total_guesses = 0;
    pthread_mutex_t mutex; // 互斥锁声明

    // 构造函数中初始化互斥锁
    PriorityQueue() {
        pthread_mutex_init(&mutex, nullptr);
    }

    // 析构函数中销毁互斥锁
    ~PriorityQueue() {
        pthread_mutex_destroy(&mutex);
    }

    void CalProb(PT& pt); // 声明保持不变
    void init(); // 声明保持不变

    // 修改 Generate 函数声明（可选）
    void Generate(PT pt);

    void PopNext();

    struct ThreadArgs {
        PriorityQueue* pq;
        string guess;
    };

    // 静态线程函数（确保没有隐藏的 this 指针问题）
    static void* add_guess(void* arg) {
        ThreadArgs* args = static_cast<ThreadArgs*>(arg);
        pthread_mutex_lock(&args->pq->mutex);
        args->pq->guesses.push_back(args->guess);
        args->pq->total_guesses++;
        pthread_mutex_unlock(&args->pq->mutex);
        return nullptr;
    }
};